#include "stream_reassembler.hh"

#include <stdexcept>

using namespace std;

void StreamReassembler::add_to_buffer(pair<size_t, string> data) {
    auto i = _buffer.begin();
    auto e = _buffer.end();

    while ((i != e) && (check_capacity(_capacity) > 0)) {
        // must be updated every loop as data may change
        auto d_first_idx = data.first;
        auto d_last_idx = d_first_idx + data.second.size() - 1;
        auto i_first_idx = i->first;
        auto i_last_idx = i_first_idx + i->second.size() - 1;

        // |data....
        //    |i....
        // add the data ahead of i to _buffer
        if (d_first_idx < i_first_idx) {
            // subtle part, contains several conditions
            auto size = check_capacity(min(data.second.size(), i_first_idx - d_first_idx));
            _buffer.insert(i, {data.first, data.second.substr(0, size)});
            _unassembled_num += size;
        }
        // ...data|
        // ...i|
        // get the data behind i to continue looping
        if (d_last_idx > i_last_idx) {
            // subtle part, contains several conditions
            data.first = max(i_last_idx + 1, d_first_idx);
            data.second = data.second.substr(data.first - d_first_idx);
            ++i;
        } else {
            // if data no more owns a rear part, end the loop
            break;
        }
    }
    // list reach end, insert the rest of the data
    if (i == e) {
        auto size = check_capacity(data.second.size());
        _buffer.insert(i, {data.first, data.second.substr(0, size)});
        _unassembled_num += size;
    }
}

StreamReassembler::StreamReassembler(const size_t capacity) : _output(capacity), _capacity(capacity) {}

//! \details This function accepts a substring (aka a segment) of bytes,
//! possibly out-of-order, from the logical stream, and assembles any newly
//! contiguous substrings and writes them into the output stream in order.
void StreamReassembler::push_substring(const string &data, const size_t index, const bool eof) {
    if (_unassembled_num + _output.buffer_size() > _capacity) {
        _output.set_error();
        throw runtime_error("exceed the total capacity");
    }
    if (eof) {
        _eof_idx = index + data.size();
    }
    if (index <= _desired_idx) {
        if (index + data.size() > _desired_idx) {
            // some data can be directly written to _output
            size_t accepted_num = _output.write(data.substr(_desired_idx - index));
            _desired_idx += accepted_num;
            // if _output have no enough space for data,
            // StreamReassembler have no space to cache it either
        }
    } else if(!data.empty()){
        // whole data have to be stored in the buffer
        add_to_buffer({index, data});
    }
    // Can data in the buffer be written to _output now?
    while ((!_buffer.empty()) && (_output.remaining_capacity() > 0)) {
        std::pair<size_t, std::string> front = {_buffer.front()};
        if (front.first > _desired_idx) {
            break;
        }
        _buffer.pop_front();
        _unassembled_num -= front.second.size();
        if (front.first + front.second.size() > _desired_idx) {
            // the parameter of _output.write() can certainly be wholely written to _output
            size_t accepted_num = _output.write(front.second.substr(_desired_idx - front.first));
            _desired_idx += accepted_num;
        }
    }
    if (_desired_idx == _eof_idx) {
        _output.end_input();
    }
}

size_t StreamReassembler::unassembled_bytes() const { return _unassembled_num; }

bool StreamReassembler::empty() const { return _unassembled_num == 0; }
